from omegaconf import OmegaConf
from torch.utils.data import DataLoader
from detectron2.config import LazyCall as L
from torch.utils.data.distributed import DistributedSampler

from data import ImageFileTrain, DataGenerator

import os

adobe_dataset_root = os.environ['ADOBE_DATASET']
coco_dataset_root = os.environ['COCO_DATASET']

# Dataloader
train_dataset = DataGenerator(
    data=ImageFileTrain(
        alpha_dir=os.path.join(adobe_dataset_root, 'full_dataset/train/alpha'),
        fg_dir=os.path.join(adobe_dataset_root, 'full_dataset/train/fg'),
        bg_dir=os.path.join(coco_dataset_root, '2014_unzip/train2014'),
        root='path/to/Adobe_Image_Matting'
    ),
    phase='train'
)

dataloader = OmegaConf.create()
dataloader.train = L(DataLoader)(
    dataset=train_dataset,
    batch_size=15,
    shuffle=False,
    num_workers=4,
    pin_memory=True,
    sampler=L(DistributedSampler)(
        dataset=train_dataset,
    ),
    drop_last=True
)
